from scipy import polyval, polyfit,  stats, randn
from pylab import plot, show
import numpy as N


def linear_regression(xnumpy,  ynumpy, polyOrder = 1):
    (m, b) = polyfit(xnumpy, ynumpy, polyOrder)
    yfit = polyval([m, b], xnumpy)
    rsq = stats.stats.pearsonr(ynumpy,yfit)
    print rsq
    return (yfit, rsq)
    
    

if __name__ == "__main__":
    x = N.arange(0, 10, 0.1)
    nse = 2 * randn(len(x))
    y2fit = 3*x+nse
    
    yfit = linear_regression(x, y2fit)
    plot(x, y2fit, 'bo')
    plot(x, yfit[0], 'r')
    show()
    print yfit[1]


    
